import numpy as np
import data
from multiprocessing import pool
import multiprocessing


def distance(x):
    a, B = x
    arr = [np.linalg.norm(a - b) for b in B]
    if len(arr) % 1000 == 0:
        print(len(arr))
    return max(arr)


def max_distance(data):
    def all_pair(data, n_thread):
        with multiprocessing.Pool(n_thread) as pool:
            res_arr = list(pool.imap(distance, zip(data, [data[i:, :] for i in range(len(data))])), )
            return res_arr

    res = all_pair(data, multiprocessing.cpu_count() // 4)
    return res


def count_distance(data_dir):
    base, d = data.fvecs_read(data_dir)
    centroid = base.mean(axis=0)
    nearest_neighbor = data.groundtruth(base, np.array([centroid]), 1)
    centroid_dis = np.linalg.norm(centroid - base[nearest_neighbor])
    max_dist = max_distance(base)
    return centroid_dis, max(max_dist)


if __name__ == "__main__":
    base_dir = '/home/zhengbian/graph-gradient/data/deepsmall_10/base.fvecs'
    centroid_distance, data_max_distance = count_distance(base_dir)
    print("deepsmall distance-nearest neighbor-centoid %.3f, dataset-max-distance %.3f" % (
        centroid_distance, data_max_distance))

    base_dir = '/home/zhengbian/graph-gradient/data/gistsmall_10/base.fvecs'
    centroid_distance, data_max_distance = count_distance(base_dir)
    print("gistsmall distance-nearest neighbor-centoid %.3f, dataset-max-distance %.3f" % (
        centroid_distance, data_max_distance))

    base_dir = '/home/zhengbian/graph-gradient/data/glovesmall_10/base.fvecs'
    centroid_distance, data_max_distance = count_distance(base_dir)
    print("glovesmall distance-nearest neighbor-centoid %.3f, dataset-max-distance %.3f" % (
        centroid_distance, data_max_distance))

    base_dir = '/home/zhengbian/graph-gradient/data/siftsmall_10/base.fvecs'
    centroid_distance, data_max_distance = count_distance(base_dir)
    print("siftsmall distance-nearest neighbor-centoid %.3f, dataset-max-distance %.3f" % (
        centroid_distance, data_max_distance))
